0373. 查找和最小的 K 对数字【中等】

• 1. 📝 题目描述
• 2. 🎯 s.1 - 最小堆

1. 📝 题目描述

• leetcode

给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。

请找到和最小的 k 个数对 (u1,v1), (u2,v2) ... (uk,vk)。

---

示例 1：

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

示例 2：

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

---

提示：

• 1 <= nums1.length, nums2.length <= 10^5
• -10^9 <= nums1[i], nums2[i] <= 10^9
• nums1 和 nums2 均为 升序排列
• 1 <= k <= 10^4
• k <= nums1.length * nums2.length

2. 🎯 s.1 - 最小堆

c

typedef struct { int sum; int i; int j; } HeapNode;

void heapSwap(HeapNode* a, HeapNode* b) { HeapNode t = *a; *a = *b; *b = t; }

void heapPush(HeapNode* heap, int* size, HeapNode node) {
    heap[(*size)++] = node;
    int i = *size - 1;
    while (i > 0 && heap[(i-1)/2].sum > heap[i].sum) {
        heapSwap(&heap[(i-1)/2], &heap[i]);
        i = (i-1)/2;
    }
}

HeapNode heapPop(HeapNode* heap, int* size) {
    HeapNode top = heap[0];
    heap[0] = heap[--(*size)];
    int i = 0;
    while (1) {
        int min = i, l = 2*i+1, r = 2*i+2;
        if (l < *size && heap[l].sum < heap[min].sum) min = l;
        if (r < *size && heap[r].sum < heap[min].sum) min = r;
        if (min == i) break;
        heapSwap(&heap[i], &heap[min]);
        i = min;
    }
    return top;
}

int** kSmallestPairs(int* nums1, int nums1Size, int* nums2, int nums2Size, int k, int* returnSize, int** returnColumnSizes) {
    *returnSize = 0;
    int** res = (int**)malloc(sizeof(int*) * k);
    *returnColumnSizes = (int*)malloc(sizeof(int) * k);
    if (nums1Size == 0 || nums2Size == 0) return res;
    int heapSize = 0;
    int initSize = nums1Size < k ? nums1Size : k;
    HeapNode* heap = (HeapNode*)malloc(sizeof(HeapNode) * (initSize + k));
    for (int i = 0; i < initSize; i++) {
        heapPush(heap, &heapSize, (HeapNode){nums1[i] + nums2[0], i, 0});
    }
    while (heapSize > 0 && *returnSize < k) {
        HeapNode node = heapPop(heap, &heapSize);
        res[*returnSize] = (int*)malloc(sizeof(int) * 2);
        res[*returnSize][0] = nums1[node.i];
        res[*returnSize][1] = nums2[node.j];
        (*returnColumnSizes)[*returnSize] = 2;
        (*returnSize)++;
        if (node.j + 1 < nums2Size)
            heapPush(heap, &heapSize, (HeapNode){nums1[node.i] + nums2[node.j+1], node.i, node.j+1});
    }
    free(heap);
    return res;
}

js

/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number} k
 * @return {number[][]}
 */
var kSmallestPairs = function (nums1, nums2, k) {
  const res = []
  // 最小堆：[sum, i, j]
  const heap = []
  const push = (val) => {
    heap.push(val)
    let i = heap.length - 1
    while (i > 0) {
      const p = (i - 1) >> 1
      if (heap[p][0] <= heap[i][0]) break
      ;[heap[p], heap[i]] = [heap[i], heap[p]]
      i = p
    }
  }
  const pop = () => {
    const top = heap[0]
    const last = heap.pop()
    if (heap.length > 0) {
      heap[0] = last
      let i = 0
      while (true) {
        let min = i
        const l = 2 * i + 1,
          r = 2 * i + 2
        if (l < heap.length && heap[l][0] < heap[min][0]) min = l
        if (r < heap.length && heap[r][0] < heap[min][0]) min = r
        if (min === i) break
        ;[heap[i], heap[min]] = [heap[min], heap[i]]
        i = min
      }
    }
    return top
  }
  for (let i = 0; i < Math.min(nums1.length, k); i++) {
    push([nums1[i] + nums2[0], i, 0])
  }
  while (heap.length > 0 && res.length < k) {
    const [, i, j] = pop()
    res.push([nums1[i], nums2[j]])
    if (j + 1 < nums2.length) push([nums1[i] + nums2[j + 1], i, j + 1])
  }
  return res
}

py

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        import heapq
        res = []
        heap = [(nums1[i] + nums2[0], i, 0) for i in range(min(len(nums1), k))]
        heapq.heapify(heap)
        while heap and len(res) < k:
            _, i, j = heapq.heappop(heap)
            res.append([nums1[i], nums2[j]])
            if j + 1 < len(nums2):
                heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
        return res

• 时间复杂度：$O(k\log k)$
• 空间复杂度：$O(k)$

算法思路：

• 初始将 $(nums1[i],nums2[0])$ 入堆，$i\in [0,min(n,k))$
• 每次弹出堆顶最小对 $(i,j)$ 后，将 $(i,j+1)$ 入堆
• 重复 $k$ 次得到结果